package org.example.myleet.p502;

import java.util.*;

/**
 * 思路：贪心+堆，首先通过思考证明贪心是可行的，以及贪心什么。
 * 贪心：考虑到第k轮选项目时，当资本达到w时，任何capital小于w的profit都可以为我所用，因此直接挑profit最大的项目做可以得到尽量大的纯利润
 * 因此这里也考虑到排序问题，那就用堆维护一个有序的候选项目profit数组
 * 堆：一个倒序的堆记录倒序的利润数组，方便每一轮选取项目
 * 做法：
 * 1、先维护好所有capital对应的profit
 * 2、然后开始模拟每一轮项目挑选
 * 3、每一轮开始时，将w可以选择的profit都加入到堆中
 * 4、然后从堆顶弹出一个profit叠加到w中
 * 5、循环3～5步，直至完成k轮项目选择
 * 6、输出最终的w
 */
public class Solution {
    public int findMaximizedCapital(int k, int w, int[] profits, int[] capital) {
        int n = profits.length;
        //先整理相同capital的profit
        Map<Integer, List<Integer>> capitalProfitsMap = new HashMap<>(n);
        for (int i = 0; i < n; ++i) {
            List<Integer> profitList = capitalProfitsMap.getOrDefault(capital[i], new LinkedList<>());
            profitList.add(profits[i]);
            capitalProfitsMap.put(capital[i], profitList);
        }
        //对capital进行排序，方便用二分法找到对应的最大下标
        Arrays.sort(capital);
        //倒序堆，存储当前可以选择的最大profit的项目
        PriorityQueue<Integer> heap = new PriorityQueue<>(Comparator.reverseOrder());
        //轮次，记录什么时候结束选择
        int epoch = 0;
        //当前已加入到堆中的capital的最大下标
        int maxIdx = -1;
        while (epoch < k) {
            //二分法找到当前w可以选取的capital的最大下标
            int idx = findCeilingCapital(capital, w);
            if (-1 == idx) {
                //选不到，就是没得选，结束吧
                break;
            }
            for (int ci = idx; ci > maxIdx; --ci) {
                //从w最新的下标开始，一直寻找到上一轮最大下标，将所有profit加入到堆中
                List<Integer> profitList = capitalProfitsMap.get(capital[ci]);
                if (profitList != null) {
                    for (Integer profit : profitList) {
                        heap.offer(profit);
                    }
                }
                //加完要扔掉，防止重复加
                capitalProfitsMap.remove(capital[ci]);
            }
            //更新已经加入堆的capital的最大下标
            maxIdx = idx;
            if (heap.size() > 0) {
                //从堆中取出一个最大利润的项目，叠加到w中
                w += heap.poll();
            }
            ++epoch;
        }
        return w;
    }

    private int findCeilingCapital(int[] capital, int target) {
        int l = 0, r = capital.length - 1;
        while (r - l > 1) {
            int m = (r + l) / 2;
            if (capital[m] > target) {
                r = m;
            } else if (capital[m] < target) {
                l = m;
            } else {
                while (m < capital.length - 1 && capital[m + 1] == target) {
                    ++m;
                }
                return m;
            }
        }
        if (capital[r] <= target) {
            while (r < capital.length - 1 && capital[r + 1] == target) {
                ++r;
            }
            return r;
        }
        if (capital[l] <= target) {
            while (l < capital.length - 1 && capital[l + 1] == target) {
                ++l;
            }
            return l;
        }
        return -1;
    }
}
